#!/usr/bin/env python3
"""
边线可视化工具（坐标映射修正版）：正确处理裁剪→缩放→边线提取的坐标映射
用法：python tools\edge_visualize.py --image assets\input.png --edges assets\edges_debug.csv --crop-csv assets\adaptive_crop.csv --output assets\edges_visualized.png
"""
import argparse
import csv
from PIL import Image, ImageDraw
import numpy as np


def load_edges_csv(csv_path):
    """加载边线数据CSV文件"""
    rows = []
    try:
        with open(csv_path, 'r', newline='') as f:
            reader = csv.DictReader(f)
            for row in reader:
                rows.append({
                    'row': int(row['row']),
                    'leftEdge': int(row['leftEdge']),
                    'rightEdge': int(row['rightEdge'])
                })
    except FileNotFoundError:
        print(f"错误：找不到边线文件 {csv_path}")
        return None
    except Exception as e:
        print(f"错误读取边线CSV文件：{e}")
        return None
    return rows


def load_crop_csv(csv_path):
    """加载裁剪信息CSV文件"""
    try:
        with open(csv_path, 'r', newline='') as f:
            reader = csv.DictReader(f)
            for row in reader:
                return {
                    'top': int(row['top']),
                    'bottom': int(row['bottom'])
                }
    except FileNotFoundError:
        print(f"警告：找不到裁剪文件 {csv_path}，将使用默认映射")
        return None
    except Exception as e:
        print(f"警告读取裁剪CSV文件：{e}，将使用默认映射")
        return None


def map_coordinates_correctly(edge_row, left_x, right_x, crop_info, 
                            original_height=480, original_width=640, 
                            scaled_height=240, scaled_width=320):
    """将缩放后的边线坐标正确映射回原始图像坐标
    
    处理流程：原始图像(480x640) → 裁剪 → 缩放(240x320) → 边线提取
    需要逆向映射：240x320 → 裁剪区域 → 480x640
    """
    
    if crop_info:
        # 有裁剪信息的情况 - 真实处理流程
        crop_top = crop_info['top']
        crop_bottom = crop_info['bottom']
        crop_height = crop_bottom - crop_top
        
        # 计算缩放比例
        scale_y = crop_height / scaled_height  # 裁剪区域高度 / 240
        scale_x = original_width / scaled_width  # 640 / 320 = 2.0
        
        # 映射Y坐标：从缩放空间回到裁剪区域，再加上裁剪偏移
        original_row = int(crop_top + edge_row * scale_y)
        
        # 映射X坐标：直接按比例缩放（320→640）
        original_left_x = int(left_x * scale_x) if left_x >= 0 else -1
        original_right_x = int(right_x * scale_x) if right_x >= 0 else -1
        
    else:
        # 无裁剪信息的默认映射（全图缩放 480×640 → 240×320）
        scale_y = original_height / scaled_height  # 480 / 240 = 2.0
        scale_x = original_width / scaled_width   # 640 / 320 = 2.0
        
        original_row = int(edge_row * scale_y)
        original_left_x = int(left_x * scale_x) if left_x >= 0 else -1
        original_right_x = int(right_x * scale_x) if right_x >= 0 else -1
    
    return original_row, original_left_x, original_right_x


def visualize_edges_corrected(image_path, edges_data, crop_info, output_path, line_width=2):
    """在原始图像上正确可视化边线（考虑完整的坐标映射）"""
    try:
        # 加载原始图像
        img = Image.open(image_path).convert('RGB')
        draw = ImageDraw.Draw(img)
        
        width, height = img.size
        print(f"原始图像尺寸：{width}x{height}")
        print(f"边线数据点数：{len(edges_data)}")
        
        # 处理流程说明
        print("坐标映射：原始图像(480x640) → 裁剪 → 缩放(240x320) → 边线提取 → 映射回原始坐标")
        
        if crop_info:
            crop_top = crop_info['top']
            crop_bottom = crop_info['bottom']
            crop_height = crop_bottom - crop_top
            print(f"裁剪区域：行 {crop_top} - {crop_bottom} (高度 {crop_height})")
            
            # 在图像上标记裁剪区域（蓝色半透明框）
            overlay = Image.new('RGBA', (width, height), (0, 0, 0, 0))
            draw_overlay = ImageDraw.Draw(overlay)
            draw_overlay.rectangle([(0, crop_top), (width-1, crop_bottom)], 
                                  fill=(100, 100, 255, 30), outline=(100, 100, 255), width=2)
            img = Image.alpha_composite(img.convert('RGBA'), overlay).convert('RGB')
            draw = ImageDraw.Draw(img)
            draw.text((10, crop_top+5), f"Crop: {crop_top}-{crop_bottom} (H={crop_height})", fill=(100, 100, 255))
        else:
            print("使用默认映射：全图缩放 480×640 → 240×320")
        
        # 统计有效边界点
        valid_left_points = []
        valid_right_points = []
        
        # 绘制左右边界（正确的坐标映射）
        for edge_data in edges_data:
            edge_row = edge_data['row']      # 缩放后的行号
            left_edge = edge_data['leftEdge']   # 缩放后的X坐标
            right_edge = edge_data['rightEdge'] # 缩放后的X坐标

            # 映射回原始图像坐标
            # 使用实际缩放高度（edges_data 的行数）代替默认 240
            scaled_h = len(edges_data)
            scaled_w = 320
            original_row, original_left_x, original_right_x = map_coordinates_correctly(
                edge_row, left_edge, right_edge, crop_info, height, width, scaled_h, scaled_w
            )
            
            # 检查映射后的坐标是否在图像范围内
            if original_row < 0 or original_row >= height:
                continue
                
            # 绘制左边界（红色）
            if original_left_x >= 0 and original_left_x < width:
                valid_left_points.append((original_left_x, original_row))
                # 绘制垂直线段增强可见性
                draw.line([(original_left_x, max(0, original_row-1)), 
                          (original_left_x, min(height-1, original_row+1))], 
                         fill=(255, 0, 0), width=line_width)
            
            # 绘制右边界（绿色）  
            if original_right_x >= 0 and original_right_x < width:
                valid_right_points.append((original_right_x, original_row))
                # 绘制垂直线段增强可见性
                draw.line([(original_right_x, max(0, original_row-1)), 
                          (original_right_x, min(height-1, original_row+1))], 
                         fill=(0, 255, 0), width=line_width)
        
        print(f"有效左边界点：{len(valid_left_points)}")
        print(f"有效右边界点：{len(valid_right_points)}")
        
        # 绘制连续边界趋势线（如果点足够多）
        if len(valid_left_points) > 10:
            # 按行号排序
            valid_left_points.sort(key=lambda p: p[1])
            # 绘制连续线（半透明红色）
            draw.line(valid_left_points, fill=(255, 150, 150), width=1)
            
        if len(valid_right_points) > 10:
            # 按行号排序  
            valid_right_points.sort(key=lambda p: p[1])
            # 绘制连续线（半透明绿色）
            draw.line(valid_right_points, fill=(150, 255, 150), width=1)
        
        # 添加图例说明
        legend_y = 20
        draw.rectangle([(10, legend_y), (25, legend_y+15)], fill=(255, 0, 0))
        draw.text((30, legend_y+2), f"Left Edge ({len(valid_left_points)} pts)", fill=(255, 0, 0))
        
        draw.rectangle([(200, legend_y), (215, legend_y+15)], fill=(0, 255, 0))
        draw.text((220, legend_y+2), f"Right Edge ({len(valid_right_points)} pts)", fill=(0, 255, 0))
        
        if crop_info:
            draw.rectangle([(400, legend_y), (415, legend_y+15)], outline=(100, 100, 255), width=2)
            draw.text((420, legend_y+2), "Crop Region", fill=(100, 100, 255))
        
        # 保存可视化结果
        img.save(output_path)
        print(f"\n✓ 边线可视化结果已保存到：{output_path}")
        
        # 输出统计信息
        if valid_left_points:
            left_x_values = [p[0] for p in valid_left_points]
            print(f"左边界X坐标范围：{min(left_x_values)} - {max(left_x_values)}")
        
        if valid_right_points:
            right_x_values = [p[0] for p in valid_right_points]
            print(f"右边界X坐标范围：{min(right_x_values)} - {max(right_x_values)}")
            
        return True
        
    except FileNotFoundError:
        print(f"错误：找不到图像文件 {image_path}")
        return False
    except Exception as e:
        print(f"可视化过程出错：{e}")
        return False


def main():
    parser = argparse.ArgumentParser(description='边线可视化工具（坐标映射修正版）')
    parser.add_argument('--image', required=True, help='输入图像路径')
    parser.add_argument('--edges', required=True, help='边线CSV文件路径')
    parser.add_argument('--crop-csv', help='裁剪信息CSV文件路径（可选）')
    parser.add_argument('--output', required=True, help='输出图像路径')
    parser.add_argument('--line-width', type=int, default=2, help='边线宽度')
    
    args = parser.parse_args()
    
    print(f"加载边线数据：{args.edges}")
    edges_data = load_edges_csv(args.edges)
    if not edges_data:
        return 1
    
    # 加载裁剪信息（如果提供）
    crop_info = None
    if args.crop_csv:
        print(f"加载裁剪信息：{args.crop_csv}")
        crop_info = load_crop_csv(args.crop_csv)
    
    print(f"可视化边线：{args.image} -> {args.output}")
    success = visualize_edges_corrected(args.image, edges_data, crop_info, args.output, args.line_width)
    
    return 0 if success else 1


if __name__ == '__main__':
    exit(main())